library(librarian)
shelf(ggpubr,
      patchwork,
      ggrepel,
      tidyplots,
      tidyverse)

theme_jpub <- function(theme_func = theme_bw) {
  theme_func(base_size = 6, base_family = 'ArialMT') +
    theme(plot.margin = margin(), legend.box.margin = margin(),
          plot.title.position = 'plot', legend.key.size = unit(4, 'mm'))
}
  
geom_dotgraph <- geom_dotplot(binaxis = 'y', stackdir = 'center')

prepare_volc <- function(df, group1 = 'group1', group2 = 'group2', fc_thres = 1, pval_thres = .05,
                       highlights = NULL, force = FALSE, auto_n = 5){
  if(is.null(df$gene)){df <- as_tibble(df, rownames = 'gene')}
  
  df <- df |>
    mutate(type = case_when(avg_log2FC > fc_thres & p_val_adj < .05 ~ group1,
                            avg_log2FC < -fc_thres & p_val_adj < .05 ~ group2,
                            .default = 'NS') |> fct_relevel(group2, 'NS'))
  
  symm_x_lim <- max(df$avg_log2FC |> abs() |> max())
  
  df |>
    dplyr::count(type) |> print()
  
  if(!is.null(highlights)){
    seek_name <- df |>
      filter(gene %in% highlights)
    if (!force) {
      seek_name <- seek_name |>
        filter(type != 'NS')
    }
  } else {seek_name <- tibble()}
  
  if(nrow(seek_name) == 0){
    warning('No interested genes upregulated or specified. Will highlight top-DEGs.')
    top.fc <- df |>
      filter(type != 'NS') |>
      group_by(type) |>
      mutate(abs_fc = abs(avg_log2FC)) |>
      slice_max(abs_fc, n = auto_n, with_ties = F) |>
      pull(gene)
    
    top.sig <- df |>
      filter(type != 'NS') |>
      group_by(type) |>
      mutate(abs_fc = abs(avg_log2FC)) |>
      slice_min(p_val_adj, n = auto_n, with_ties = F) |>
      pull(gene)
    
    seek_name <- df |>
      filter(gene %in% c(top.sig, top.fc))
  }
  
  list(df, seek_name, symm_x_lim)
}

#' Plot volcano for generic & slides 
#'
#' @param df data.frame from Seurat::FindMarker() output.
#' @param group1 character vector of ident.1 name, for legend.
#' @param fc_thres log2fc threshold for dash line threshold, 1 by default.
#' @param pval_thres adjusted p value threshold for dash line threshold, 0.05 by default.
#' @param highlights optional character string of gene names wanted to highlight. Will auto highlight genes if leave empty.
#' @param force Boolean, highlight `highlights` genes anyway, even they are not significantly upregulated? FALSE by default.
#'
#' @return A ggolot object.
#' @export
#'
#' @examples 
#' pbmc_small |>
#'   FindMarker(ident.1 = 0) |>
#'   plot_bill_volc('clus0')
#'   
plot_bill_volc <- function(df, group1, group2, fc_thres = 1, pval_thres = .05,
                          highlights = NULL, force = FALSE){
  df_list <- prepare_volc(df, group1, group2, fc_thres, pval_thres,
               highlights, force)
  seek_name <- df_list[[2]]
  symm_x_lim <- df_list[[3]]
  
  df_list[[1]] |>
    ggplot(aes(avg_log2FC, -log10(p_val_adj), color = type)) +
    geom_point(size = .5, alpha = .3) +
    geom_vline(xintercept = c(-fc_thres,fc_thres), linetype = 'dashed') +
    geom_hline(yintercept = -log10(pval_thres), linetype = 'dashed') +
    geom_point(data = seek_name, color = 'purple') +
    geom_text_repel(data = seek_name, aes(label = gene), color = 'black') +
    scale_color_manual(values = c('royalblue','grey','red2') ) +
    theme_pubr() +
    theme(legend.position = 'top', plot.margin = margin(), legend.box.margin = margin()) +
    expand_limits(x = c(-symm_x_lim,symm_x_lim))
}

# TODO: refactor two volc func to avoid replicate self! =========

#' Plot volcano for journal-publication 
#'
#' @param df data.frame from Seurat::FindMarker() output.
#' @param group1 character vector of ident.1 name, for legend.
#' @param fc_thres log2fc threshold for dash line threshold, 1 by default.
#' @param pval_thres adjusted p value threshold for dash line threshold, 0.05 by default.
#' @param highlights optional character string of gene names wanted to highlight. Will auto highlight genes if leave empty.
#' @param force Boolean, highlight `highlights` genes anyway, even they are not significantly upregulated? FALSE by default.
#'
#' @return A ggolot object.
#' @export
#'
#' @examples 
#' pbmc_small |>
#'   FindMarker(ident.1 = 0) |>
#'   plot_pub_volc('clus0')
#'   
plot_pub_volc <- function(df, group1, group2, fc_thres = 1, pval_thres = .05,
                           highlights = NULL, high_color = 'purple', force = FALSE, ...){
  df_list <- prepare_volc(df, group1, group2, fc_thres, pval_thres,
                          highlights, force, ...)
  seek_name <- df_list[[2]]
  symm_x_lim <- df_list[[3]]
  
  df_list[[1]] |>
    ggplot(aes(avg_log2FC, -log10(p_val_adj), color = type)) +
    geom_point(size = .3, alpha = .3) +
    geom_vline(xintercept = c(-1,1), linetype = 'dashed', alpha = .5) +
    geom_hline(yintercept = 1.3, linetype = 'dashed', alpha = .5) +
    geom_point(data = seek_name, size = 1, color = high_color) +
    geom_text_repel(data = seek_name, aes(label = gene), color = 'black',
                    size = 2, box.padding = .1) +
    scale_color_manual(values = c('royalblue','grey','red2') ) +
    theme_classic(base_size = 6) +
    theme(legend.position = 'top', plot.margin = margin(), legend.box.margin = margin(),
          legend.key.size = unit(4, 'mm')) +
    expand_limits(x = c(-symm_x_lim,symm_x_lim))
}

plot_enrichment <- function(df, base_col = 'Reds', n = 10, descr_wrap = 40,
                            metric = NULL, padj_thres = .05, force_regex = NULL){
  if(!is.data.frame(df)) {
    df <- df@result
  }
  if(is.null(metric)) {
    metric <- case_when('Count' %in% colnames(df) ~ 'Count',
                        'NES' %in% colnames(df) ~ 'NES',
                        .default = NA)
    if(is.na(metric)) stop('Count and NES not found in `df`. Please supply a proper `metric`.')
  }
  if (!is.null(force_regex)) {
    highlight <- df |>
      filter(str_detect(Description, force_regex))
    if (nrow(highlight) == 0) {
      warning('Failed to find highlighted pathway!')
    }
    plot_df <- df |>
      filter(p.adjust < padj_thres) |>
      slice_sample(n = n-1) |>
      bind_rows(highlight) |>
      mutate(Description = str_wrap(Description, descr_wrap) |>
               fct_reorder(.data[[metric]]))
  } else {
    plot_df <- df |>
      filter(p.adjust < padj_thres) |>
      slice_min(p.adjust, n = n, with_ties = F) |>
      mutate(Description = str_wrap(Description, descr_wrap) |>
               fct_reorder(.data[[metric]]))
  }
  plot_df |>
    ggplot(aes(x = .data[[metric]], y = Description, fill = p.adjust)) +
    geom_col() +
    theme_pubr(legend = 'right') +
    labs_pubr() +
    scale_fill_distiller(palette = base_col)
}

#' Calculate fraction confidence interval of subtype in group
#' 
#' Useful for ggplot col plot with error bars.
#'
#' @param df A tibble with `group` and `subtype` columns, typically a scRNA-seq meta data.
#' @param group colname of column containing variable of interest (e.g. WT/KO), need to have more than 1 types.
#' @param subtype colname of column containing variable of cell types (e.g. CD4T/CD8T/NK etc.)
#'
#' @return A tibble added some columns: 
#' - `fraction`: subtype fraction in group
#' - `conf.low` & `conf.high`: 95% CI of fraction, by `prop.test`
#' - `group.sum`: sum of group count
#' @export
#'
#' @examples 
#' meta_kera |>
#'   calc_frac_conf_on_grouped_count(orig.ident, seurat_clusters)
#' 
calc_frac_conf_on_grouped_count <- function(df, group, subtype){
  df |>
    dplyr::count({{ group }}, {{ subtype }}) |>
    group_by({{ group }}) |>
    mutate(group.sum = sum(n),
           fraction = n / group.sum) |>
    rowwise() |>
    mutate(conf.low = prop.test(n, group.sum, correct = F)$conf.int[1],
           conf.high = prop.test(n, group.sum, correct = F)$conf.int[2]) |>
    ungroup()
}


#' Test significance for fraction of subtype in group
#'
#' @param df A tibble with `group` and `subtype` columns, typically a scRNA-seq meta data.
#' @param group colname of column containing variable of interest (e.g. WT/KO), need to have 2 and only 2 types.
#' @param subtype colname of column containing variable of cell types (e.g. CD4T/CD8T/NK etc.)
#'
#' @return A tibble with 2 columns: 
#' - `subtype`: defined `subtype`
#' - `p.value`: p value calculated from `fisher.test`
#' @export
#'
#' @examples
#' meta_kera |>
#'   test_on_grouped_count(orig.ident, seurat_clusters)
#'   
test_on_grouped_count <- function(df, group, subtype){
  df |>
    dplyr::count({{ group }}, {{ subtype }}) |>
    group_by({{ group }}) |>
    mutate(n0 = sum(n) - n) |>
    ungroup() |>
    pivot_wider(names_from = {{ group }}, values_from = c(n, n0)) |>
    na.omit() |>
    set_names(c('subtype', 'n1_x', 'n1_y', 'n0_x', 'n0_y')) |>
    rowwise() |>
    mutate(p.value = fisher.test(
      matrix(c(n1_x, n1_y, n0_x, n0_y), ncol = 2))$p.value) |>
    select(subtype, p.value) |>
    ungroup()
}

#' Compare cell type fraction changes between 2 conditions
#'
#' @param df A tibble with `group` and `subtype` columns, typically a scRNA-seq meta data.
#' @param group colname of column containing variable of interest (e.g. WT/KO), need to have 2 and only 2 types.
#' @param subtype colname of column containing variable of cell types (e.g. CD4T/CD8T/NK etc.)
#' @param var.1 identity class in `group` as numerator to calculate log2FC.
#' @param var.2 identity class in `group` as denominator for comparison.
#'
#' @return A tibble with 6 columns: 
#' - `subtype`: defined `subtype`
#' - `p.value`: p value calculated from t.test
#' - 2 columns named the same as `var.1` and `var.2`: their fraction value
#' - `log2fc_frac`: log2 of fold change of `var.1`/`var.2`
#' - `type`: 'NS''Decreased''Increased' for adjusted p value in `var.1`/`var.2`
#' @export
#'
#' @examples
#' meta_kera |>
#'   discov_frac_change(orig.ident, seurat_clusters)
#'   
discov_frac_change <- function(meta, group, subtype, var.1, var.2) {
  if (isS4(meta)) {
    meta <- meta@meta.data |> as_tibble(rownames = '.cell')
  }
  
  fmeta <- meta |>
    mutate(fvar = {{ group }}, ftype = {{ subtype }}, .keep = 'none')
  
  logfc_x <- fmeta |>
    calc_frac_conf_on_grouped_count(fvar, ftype) |>
    pivot_wider(names_from = fvar, values_from = fraction,
                id_cols = ftype) |>
    mutate(log2fc_frac = log2({{ var.1 }}/{{ var.2 }}))
  
  fmeta |>
    test_on_grouped_count(fvar, ftype) |>
    left_join(logfc_x, join_by(subtype == ftype)) |>
    mutate(p.value = p.adjust(p.value, method = 'bonferroni'),
           type = case_when(
             p.value < .05 & log2fc_frac > 0 ~ 'Increased',
             p.value < .05 & log2fc_frac < 0 ~ 'Decreased',
             .default = 'NS'),
           subtype = fct_reorder(subtype, log2fc_frac))
}


publish_pdf <- function(file, project = proj.nm, width = 50, height = 50){
  if (is.character(project)) {
    file = str_glue('{project}/figs/{file}')
  }
  ggsave(file, width = width, height = height, units = 'mm')
}

publish_source_plot <- function(prefix, project = proj.nm, width = 50, height = 50,
                                g0 = last_plot()) {
  if (is(g0, 'patchwork')) {
    message('INFO: last_plot() return a patchwork object.')
    if (length(g0) > 1) {
      warning('More than one object in patchwork. Save only first object.')
    }
    g0 <- g0[[1]]
  }
  plot_file = str_glue('{project}/figs/{prefix}.pdf')
  ggsave(plot_file, plot = g0, width = width, height = height, units = 'mm')
  g0 |>
    pluck('data') |>
    write_source_csv(prefix = prefix, project = project)
}

write_source_csv <- function(df, prefix, project = proj.nm, ...) {
  res_dir <- str_c(project, '/results/')
  csv_path <- str_c(res_dir, prefix, '.csv')
  if (!dir.exists(res_dir)) {
    dir.create(res_dir, mode = '0775')
    message('"results" dir not found in project dir. Create it...')
  }
  message(str_glue('Write data to {csv_path}...'))
  write_csv(df, csv_path, ...)
}


#' Make dataframe `x` for ggpubr::stat_pvalue_manual(data = x)
#'
#' @param df The df containing all following variables, for adding layer of stat_pvalue_manual().
#' @param y The variable name of y in ggplot(aes()).
#' @param facets Required. The facet variable in main ggplot .
#' @param p.value The variable containing p value.
#' @param group.1 Character. The group 1 in testing.
#' @param group.2 Character. The group 2 in testing.
#'
#' @return A df for ggpubr::stat_pvalue_manual(data = x)
#' @export
#'
#' @examples
customize_pvalue <- function(df, y, facets, p.value, group.1, group.2){
  df |>
    group_by({{ facets }}) |>
    reframe(y.position = max({{ y }}) * 1.05,
            p = mean({{ p.value }}) |> signif(3),
            group1 = group.1,
            group2 = group.2)
}

add_manual_pvalue <- function(x){
  stat_pvalue_manual(data = x,label = 'p',hide.ns = 'p')}

#' Plot MAF by 1KGP subpopulation of SNP
#'
#' @param df A data.frame or its extension. Must containing `maf` column and column specified in `pop_column`.
#' @param pop_column A character. Column name of population codes (like CHS).
#' @param snp A character. SNP name displayed in plot title.
#'
#' @return A ggplot object.
#' @export
#'
#' @examples
#' plot_1kgp_globe(df = mysnp.with.pop, pop_column = 'subpop', snp = 'IGHG1-G396R')
plot_1kgp_globe <- function(df, pop_column, snp){
  df <- read_delim('Archive/covid19/ref/igsr_populations.tsv',
                               name_repair = make.names, show_col_types = F) |>
    filter(!is.na(Population.code) & !is.na(Superpopulation.code)) |>
    right_join(df, join_by('Population.code' == {{pop_column}})) |>
    mutate(maf = signif(maf, 3))
  
  world <- map_data("world")
  
  world |>
    ggplot(aes(x = long, y = lat, group = group)) + 
    geom_polygon(fill = 'white', color = 'grey') +
    geom_point(data = df,
               aes(x = Population.longitude,
                   y = Population.latitude,
                   color = maf),
               size = 3,
               inherit.aes = FALSE) +
    geom_text_repel(data = filter(df, maf > 0),
                    aes(x = Population.longitude,
                        y = Population.latitude,
                        label = str_c(maf * 100, '%')),
                    nudge_y = 5,
                    inherit.aes = FALSE) +
    theme_bw() +
    scale_color_viridis_c(limits = c(0,NA)) +
    coord_cartesian(xlim = c(-160,170), ylim = c(-50,80)) +
    labs(x = 'longitude', y = 'latitude',
         color = 'Minor allele freq', title = str_glue('{snp} in 1KGP data'))}

anno_pmc_hits <- function(kw.list) {
  require(europepmc)
  map_vec(kw.list, epmc_hits, .progress = T)
}

search_go_term <- function(term) {
  require(GO.db)
  GO.db |> keys('TERM') |>
    str_subset(term) |>
    AnnotationDbi::select(x = GO.db, keys = _,
                          keytype = 'TERM', columns = 'GOID') |>
    as_tibble()
}

map_go_gene <- function(goid, org = 'human') {
  if (org == 'human') {
    require(org.Hs.eg.db)
    res <-
    AnnotationDbi::select(org.Hs.eg.db, keys = goid,
                          keytype = 'GOALL', columns = 'SYMBOL')
  } else if (org == 'mouse') {
    require(org.Mm.eg.db)
    res <-
    AnnotationDbi::select(org.Mm.eg.db, keys = goid,
                          keytype = 'GOALL', columns = 'SYMBOL')
  } else {
    stop('Please use "human" or "mouse" for `org`!')
  }
  
  res |>
    dplyr::select(GOALL, SYMBOL) |>
    distinct(GOALL, SYMBOL, .keep_all = T) |>
    as_tibble()
}

#' Query human variation info from Ensembl Rest API
#'
#' @param rsid Character. Variant identifier, including structural variant identifiers.
#' @param type Character. One of 'genotypes' (individual genotypes), 'population_genotypes' (population genotype frequencies, alias 'pop_geneotypes'), 'pops' (population allele frequencies, alias 'allele'), 'genotyping_chips', 'phenotypes'.
#'
#' @return A tibble parsed from json response.
#' @export
#'
#' @examples
#' query_snp_maf('rs1050501')
#' query_snp_maf('rs1050501', type = 'genotypes')
#' query_snp_maf('rs117518546', type = 'pop_genotypes')
#' 
query_snp_maf <- function(rsid, type = 'allele') {
  require(httr)
  
  option <- case_match(type,
                       'allele' ~ 'pops',
                       'pop_genotypes' ~ 'population_genotypes',
                       .default = type)
  
  content.slot <- case_match(type,
                             c('allele','pops') ~ 'populations',
                             'pop_genotypes' ~ 'population_genotypes',
                             .default = type)

  server <- "https://rest.ensembl.org"
  ext <- "/variation/human/{rsid}?{option}=1" |>
    str_glue()
  
  r <- GET(paste(server, ext, sep = ""), content_type("application/json"))
  
  stop_for_status(r)
  
  content(r) |>
    pluck(content.slot) |>
    map(as_tibble) |>
    list_rbind() |>
    mutate(sample = str_remove(sample, '.+:'))
}

query_uniprot_keyword <- function(qvalue, qkey = 'keyword', species = "human") {
  sp.code <- case_match(species,
                            'human' ~ '9606',
                            'mouse' ~ '10090',
                            .default = species)
  domain <- "https://rest.uniprot.org/uniprotkb/stream?compressed=true&fields=accession%2Cgene_primary&format=tsv"
  query <- "&query=((taxonomy_id:{sp.code})+AND+({qkey}:{qvalue}))+AND+(reviewed:true)" |>
    str_glue()
  str_c(domain, query) |>
    download.file(destfile = '/tmp/uniprot.keyword')
  
  read_delim('/tmp/uniprot.keyword') |>
    dplyr::rename("symbol" = "Gene Names (primary)") |>
    separate_longer_delim(symbol, delim = "; ") |>
    distinct(symbol, .keep_all = TRUE)
}

query_litvar_rsid <- function(symbol, mutation) {
  require(httr)
  
  server <- "https://www.ncbi.nlm.nih.gov/research/litvar2-api/variant"
  ext <- str_glue("/autocomplete/?query={symbol}%20{mutation}")
  
  r <- GET(paste(server, ext, sep = ""), content_type("application/json"))
  
  stop_for_status(r)
  
  content(r)[[1]][['rsid']]
}

BubblePlot <- function(df, scaled = TRUE, d2 = FALSE, size = c(0,3)) {
  aes_params <- list(x = if(d2) quote(group.x) else quote(features.plot),
                     y = if(d2) quote(group.y) else quote(id),
                     fill = if(scaled) quote(avg.exp.scaled) else quote(avg.exp))
  
  g1 <- df |>
    ggplot(aes(!!aes_params$x, !!aes_params$y, size = pct.exp, fill = !!aes_params$fill)) +
    geom_point(shape = 21, stroke = .3) +
    scale_radius(range = size, breaks = scales::breaks_extended(4)) +
    theme_pubr(legend = "right") +
    labs(fill = "Average expression", size = "Percent expressed",
         x = if(d2) 'Group X' else 'Gene',
         y = if(d2) 'Group Y' else 'Cell type')
  
  if(scaled) {
    g1 + scale_fill_distiller(palette = 'RdBu',
                              values = pretty_distiller(df$avg.exp.scaled))
  } else {
    g1 + scale_fill_distiller(palette = 'Reds', direction = 1)
  }
}

pretty_distiller <- function(x) {
  min_max <- range(x)
  midpoint <- (-min_max[1])/(min_max[2] - min_max[1])
  if (midpoint < 0 || midpoint > 1) stop('Value range should include 0!')
  c(0, midpoint, 1)
}

#' Calculate odds ratio for 4 SNP-disease models
#'
#' @param df A dataframe. It needs to have `cohort`, `genotype`, `count` columns. 
#' @param ref_group Character. Groups in `cohort` column of `df`
#' @param alt_allele Character. Allele string appeared in `genotype` column of `df`
#' @param covariate (Optional) Character. Column name in `df` containing covariates
#'
#' @return A tibble containing `or_value`, `ci_lower`, `ci_upper`, `p.value`, `model` 5 columns and 4 rows with 'Additive', 'Allelic', 'Recessive', 'Dominant' in `model` column.
#' @export
#'
#' @examples
#' foo <-
#' tibble(genotype = sample(c('CC','CT','TT'), 100, replace = T, prob = c(.1,.3,.6)),
#'        disease = sample(0:1, 100, replace = T),
#'        sex = sample(0:1, 100, replace = T)) |>
#'        summarize(count = n(), .by = everything())
#' 
#' foo |> model_OR_forests(ref_group = '0', alt_allele = 'T',
#' covariate = 'sex')
model_OR_forests <- function(df, ref_group, alt_allele, covariate = '') {
  regress_logistic <- function(model_df, model_name, covariate = covariate) {
    glm_formula <- as.formula(str_glue('disease ~ {model_name}{covariate}'))
    
    model_df |>
      glm(glm_formula, family = binomial(link = "logit"), data = _,
          weights = count) |>
      generics::tidy() |>
      filter(str_detect(term, model_name)) |>
      mutate(or_value = exp(estimate), p.value, model = model_name,
             ci_lower = exp(estimate-1.96*std.error),
             ci_upper = exp(estimate+1.96*std.error), .keep = 'none')
  }
  
  df_model3 <- df |>
    mutate(Additive = str_count(genotype, alt_allele),
           Recessive = Additive == 2,
           Dominant = Additive > 0,
           disease = cohort != ref_group)
  
  df_rec <- df_model3 |>
    regress_logistic(model_name = 'Recessive', covariate = covariate)
  
  df_dom <- df_model3 |>
    regress_logistic(model_name = 'Dominant', covariate = covariate)
  
  df_ale <- df_model3 |>
    mutate(Allelic = case_match(Additive,
                               0 ~ list(c(0,0)),
                               1 ~ list(c(0:1)),
                               2 ~ list(c(1,1)))) |>
    unnest(Allelic) |>
    regress_logistic(model_name = 'Allelic', covariate = covariate)
  
  df_add <- df_model3 |>
    regress_logistic(model_name = 'Additive', covariate = covariate)
  
  list(df_add, df_ale, df_rec, df_dom) |>
    bind_rows()
}

batch_enrich_path <- function(df, org = 'Hs', path = c('GO', 'KEGG'),
                              method = c('GSEA', 'ORA'), logfc_threshold = 1) {
  require(clusterProfiler)
  path <- match.arg(path)
  method <- match.arg(method)
  if (path == 'KEGG') {
    df <- df$gene |>
      bitr(fromType = 'SYMBOL', toType = 'ENTREZID',
           OrgDb = str_glue('org.{org}.eg.db')) |>
      mutate(gene = SYMBOL) |>
      left_join(df) |>
      mutate(gene = ENTREZID)
  }
  
  clist <- df$cluster |>
    unique() |>
    set_names() 
  
  if (method == 'ORA') {
    clist <- expand_grid(clist, c('up', 'down')) |>
      pmap_chr(paste, sep = '_') |>
      set_names()
    
    print(clist)
    
    df <- df |>
      filter(p_val_adj < .05, abs(avg_log2FC) > logfc_threshold) |>
      mutate(type = ifelse(avg_log2FC > 0, 'up', 'down'),
             cluster = str_c(cluster, type, sep = '_'))
    
    message('prepare ORA list...')
    #print(df)
    
    enrich_glist <- clist |>
      map(\(x)df |> filter(cluster == x) |>
            pull(gene))
    
    message('prepare gene list...')
    print(head(enrich_glist[[1]]))
    
    if (path == 'KEGG') {
      enrich_glist |>
        map(safely(\(x)enrichKEGG(x, organism = ifelse(org == 'Hs', 'hsa', 'mmu'))),
            .progress = T)
    } else {
      enrich_glist |>
        map(\(x)x |>
              enrichGO(ont = 'ALL', OrgDb = str_glue('org.{org}.eg.db'),
                       keyType = 'SYMBOL', readable = T), .progress = T)
    }
  } else {
    enrich_glist <- clist |>
      map(\(x)df |> filter(cluster == x, p_val_adj < .05) |>
            pull(avg_log2FC, name = gene) |>
            sort(decreasing = T))
    
    if (path == 'KEGG') {
      enrich_glist |>
        map(\(x)gseKEGG(x, organism = ifelse(org == 'Hs', 'hsa', 'mmu'),
                        eps = 0, pvalueCutoff = 1))
    } else {
      enrich_glist |>
        map(\(x)x |>
              gseGO(ont = 'ALL', OrgDb = str_glue('org.{org}.eg.db'),
                    keyType = 'SYMBOL', eps = 0, pvalueCutoff = 1))
    }}
}

#' Read .RData/.rda file into an object or list
#'
#' Reads objects from a .RData file. If the file contains a single object,
#' the object itself is returned. If the file contains multiple objects,
#' they are returned as a named list.
#'
#' @param rda_path A character string giving the path to the .RData or .rda file.
#'
#' @return The loaded object. If the file contains one object, that object is
#'   returned. If multiple, a named list is returned.
#' @export
#'
#' @examples
#' \dontrun{
#' # --- Setup: Create example .RData files ---
#' save(data.frame(x = 1:3), file = "single.RData")
#' save(a = 1, b = 2, file = "multiple.RData")
#'
#' # 1. Read a single object
#' my_df <- read_rda("single.RData")
#' print(my_df)
#'
#' # 2. Read multiple objects into a list
#' my_list <- read_rda("multiple.RData")
#' print(my_list)
#'
#' # 3. Use in a pipe
#' library(magrittr)
#' read_rda("single.RData") %>% summary()
#'
#' # Clean up
#' file.remove("single.RData", "multiple.RData")
#' }
read_rda <- function(rda_path) {
  if (!file.exists(rda_path)) {
    stop("File not found: ", rda_path)
  }
  
  temp_env <- new.env()
  loaded_names <- load(rda_path, envir = temp_env)
  
  if (length(loaded_names) == 0) {
    warning("The file '", rda_path, "' is empty or contains no R objects.")
    return(invisible(NULL))
  }
  
  if (length(loaded_names) == 1) {
    return(get(loaded_names[1], envir = temp_env))
  } else {
    return(mget(loaded_names, envir = temp_env))
  }
}

#' Title
#'
#' @param robj 
#' @param file 
#' @param level 
#' @param threads 
#'
#' @returns
#' @export
#'
#' @examples
#' bench::press(n_threads = 1:8,
#' {
#' bench::mark(
#' write_zst = write_zstd_rds(sobj, 'foo.zstd.rds', threads = n_threads)
#' )
#' })
write_zstd_rds <- function(robj, file, level = 3, threads = 4) {
  require(zstdlite)
  con <- zstdfile(file, level = level, num_threads = threads)
  on.exit(close(con))
  suppressWarnings(saveRDS(robj, file = con, compress = F))
}

read_zstd_rds <- function(file) {
  require(zstdlite)
  con <- zstdfile(file)
  on.exit(close(con))
  readRDS(file = con)
}